import numpy as np
import matplotlib.pyplot as plt
from scipy.special import erfc

# 信噪比范围
SNR_db = np.arange(-5, 12.1, 0.2)
SNR_lin = 10**(SNR_db/10)

# 理论误码率计算
def calc_ber(mod_type, snr, optimal=False):
    if mod_type == 'ASK':
        eb_n0 = snr/4 if not optimal else snr/2
        return 0.5 * erfc(np.sqrt(eb_n0))
    elif mod_type == 'PSK':
        eb_n0 = snr if not optimal else 2*snr
        return 0.5 * erfc(np.sqrt(eb_n0))
    elif mod_type == 'FSK':
        eb_n0 = snr/2 if not optimal else snr
        return 0.5 * erfc(np.sqrt(eb_n0))

# 绘制实际接收曲线
plt.figure(figsize=(10,6))
for mod in ['ASK', 'PSK', 'FSK']:
    ber = [calc_ber(mod, snr) for snr in SNR_lin]
    plt.semilogy(SNR_db, ber, label=f'2{mod} Practical Receiver', linewidth=2)
plt.xlabel('SNR (dB)')
plt.ylabel('Bit Error Rate (BER)')
plt.legend()
plt.grid(True, which='both', alpha=0.3)
plt.title('Practical Receiver Performance')
plt.ylim(1e-7, 1)
plt.xlim(-5,12)
plt.savefig('practical_ber.png', dpi=300, bbox_inches='tight')
plt.show()

# 绘制最佳接收曲线
plt.figure(figsize=(10,6))
for mod in ['ASK', 'PSK', 'FSK']:
    ber = [calc_ber(mod, snr, optimal=True) for snr in SNR_lin]
    plt.semilogy(SNR_db, ber, label=f'2{mod} Optimal Receiver', linewidth=2)
plt.xlabel('SNR (dB)')
plt.ylabel('Bit Error Rate (BER)')
plt.legend()
plt.grid(True, which='both', alpha=0.3)
plt.title('Optimal Receiver Performance')
plt.ylim(1e-7, 1)
plt.xlim(-5,12)
plt.savefig('optimal_ber.png', dpi=300, bbox_inches='tight')
plt.show()

# 对比曲线
plt.figure(figsize=(12,10))
mods = ['ASK', 'PSK', 'FSK']
for i, mod in enumerate(mods):
    plt.subplot(3,1,i+1)
    plt.semilogy(SNR_db, [calc_ber(mod,snr) for snr in SNR_lin],
                 label='Practical Receiver', linewidth=2)
    plt.semilogy(SNR_db, [calc_ber(mod,snr,optimal=True) for snr in SNR_lin],
                 label='Optimal Receiver', linewidth=2)
    plt.title(f'2{mod} Performance Comparison')
    plt.grid(True, which='both', alpha=0.3)
    plt.ylabel('BER')
    plt.legend()
    if i == 2:
        plt.xlabel('SNR (dB)')
    plt.ylim(1e-7, 1)
    plt.xlim(-5,12)

plt.tight_layout()
plt.savefig('ber_comparison.png', dpi=300, bbox_inches='tight')
plt.show()

# 所有调制方式在同一图中对比
plt.figure(figsize=(12,8))

# 实际接收
plt.subplot(2,1,1)
for mod in ['ASK', 'PSK', 'FSK']:
    ber = [calc_ber(mod, snr) for snr in SNR_lin]
    plt.semilogy(SNR_db, ber, label=f'2{mod}', linewidth=2)
plt.title('Practical Receiver Performance - All Modulations')
plt.ylabel('BER')
plt.legend()
plt.grid(True, which='both', alpha=0.3)
plt.ylim(1e-7, 1)
plt.xlim(-5,12)

# 最佳接收
plt.subplot(2,1,2)
for mod in ['ASK', 'PSK', 'FSK']:
    ber = [calc_ber(mod, snr, optimal=True) for snr in SNR_lin]
    plt.semilogy(SNR_db, ber, label=f'2{mod}', linewidth=2)
plt.title('Optimal Receiver Performance - All Modulations')
plt.xlabel('SNR (dB)')
plt.ylabel('BER')
plt.legend()
plt.grid(True, which='both', alpha=0.3)
plt.ylim(1e-7, 1)
plt.xlim(-5,12)

plt.tight_layout()
plt.savefig('all_modulations_ber.png', dpi=300, bbox_inches='tight')
plt.show()

# 计算在BER=1e-3时所需的SNR
target_ber = 1e-3
print("Required SNR for BER = 1e-3:")
for mod in ['ASK', 'PSK', 'FSK']:
    # 实际接收
    practical_ber = [calc_ber(mod, snr) for snr in SNR_lin]
    practical_idx = np.argmin(np.abs(np.array(practical_ber) - target_ber))
    practical_snr = SNR_db[practical_idx]

    # 最佳接收
    optimal_ber = [calc_ber(mod, snr, optimal=True) for snr in SNR_lin]
    optimal_idx = np.argmin(np.abs(np.array(optimal_ber) - target_ber))
    optimal_snr = SNR_db[optimal_idx]

    print(f"2{mod}: Practical = {practical_snr:.1f} dB, Optimal = {optimal_snr:.1f} dB")